﻿#pragma once

#include <array>
#include <cstddef>
#include <cstdint>
#include <stdexcept>
#include <type_traits>
#include <vector>

namespace xy
{

namespace math
{

// 定义 D 维向量参数概念
template <std::size_t D, class... Ts>
concept DimParams = requires { sizeof...(Ts) == D && (std::is_integral_v<Ts> && ...); };

template <std::size_t N, class T, std::size_t LoBound = 0, std::size_t HiBound = LoBound> class ndarray
{
  protected:
    static_assert(N > 0, "N cannot be 0");

    // Dim 索引，由于 LoBound 和 HiBound 的存在，可以为负值
    using Dim = std::array<std::intptr_t, N>;
    using Shape = std::array<std::size_t, N>;

    std::vector<T> m_arr;
    Shape m_shape{};

    /**
     * @brief 计算数组大小
     *
     * @param[in] shape
     * @return constexpr std::size_t
     */
    constexpr static std::size_t _calc_size(const Shape &shape) noexcept
    {
        std::size_t size = shape[0] + (LoBound + HiBound);
        for (std::size_t i = 1; i < N; i++)
            size *= shape[i] + (LoBound + HiBound);
        return size;
    }

  public:
    ndarray(const ndarray &) = default;
    ndarray(ndarray &&) = default;
    ndarray &operator=(const ndarray &) = default;
    ndarray &operator=(ndarray &&) = default;
    ~ndarray() = default;

    /**
     * @brief 通过形状构造数组
     *
     * @param[in] shape
     */
    explicit ndarray(const Shape &shape) : m_arr(_calc_size(shape)), m_shape(shape)
    {
    }

    /**
     * @brief 通过形状构造数组并赋值
     *
     * @param[in] shape
     */
    explicit ndarray(const T &value, const Shape &shape) : m_arr(_calc_size(shape), value), m_shape(shape)
    {
    }

    /**
     * @brief 多维形状
     *
     * @tparam Ts
     */
    template <class... Ts>
        requires(DimParams<N, Ts...>)
    explicit ndarray(const Ts &...ts) : ndarray(Shape{std::size_t(ts)...})
    {
    }

    /**
     * @brief 多维形状并赋值
     *
     * @tparam Ts
     */
    template <class... Ts>
        requires(DimParams<N, Ts...>)
    explicit ndarray(const T &value, const Ts &...ts) : ndarray(value, Shape{std::size_t(ts)...})
    {
    }

    /**
     * @brief 通过 initializer_list 构造数组
     * 注意使用此构造函数时，必须在之后调用 fit_shape() 进行形状调整
     *
     * @param[in] l
     */
    explicit ndarray(std::initializer_list<T> l)
    {
        for (auto &v : l)
            m_arr.emplace_back(std::move(v));
    }

    /**
     * @brief 通过 std::vector 构造数组
     * 注意使用此构造函数时，必须在之后调用 fit_shape() 进行形状调整
     *
     * @param[in] l
     */
    explicit ndarray(std::vector<T> &&l)
    {
        for (auto &v : l)
            m_arr.emplace_back(std::move(v));
    }

    /**
     * @brief 调整 shape 与 arr 大小一致
     *
     * @param[in] shape
     */
    void fit_shape(const Shape &shape)
    {
        std::size_t size = _calc_size(shape);
        if (m_arr.size() != size)
            throw std::logic_error("Shape not match.");
        m_shape = shape;
    }

    /**
     * @brief 调整 shape 和 arr 的大小
     *
     * @tparam Ts
     */
    template <class... Ts>
        requires(DimParams<N, Ts...>)
    void reshape(const Ts &...ts)
    {
        m_shape = Shape{std::size_t(ts)...};
        std::size_t size = _calc_size(m_shape);
        m_arr.resize(size);
    }

    /**
     * @brief 调整 shape 和 arr 的大小并赋值
     *
     * @tparam Ts
     */
    template <class... Ts>
        requires(DimParams<N, Ts...>)
    void reshape(const T &value, const Ts &...ts)
    {
        m_shape = Shape{std::size_t(ts)...};
        std::size_t size = _calc_size(m_shape);
        m_arr.resize(size, value);
    }

    /**
     * @brief 获得第 i 维度大小
     *
     * @param[in] i
     * @return constexpr std::size_t
     */
    constexpr std::size_t shape(std::size_t i) const noexcept
    {
        return m_shape[i];
    }

    /**
     * @brief 获得数组大小
     *
     * @return constexpr std::size_t
     */
    constexpr std::size_t size() const noexcept
    {
        return m_arr.size();
    }

    /**
     * @brief 根据 dim 获得对应的索引（注意第一个维度是最高维度）
     *
     * @param[in] dim
     * @return constexpr std::size_t
     */
    constexpr std::size_t index(const Dim &dim) const noexcept
    {
        std::size_t offset{dim[N - 1] + LoBound};
        std::size_t term = 1;
        for (std::size_t i = 1; i < N; i++)
        {
            term *= m_shape[N - i] + (LoBound + HiBound);
            offset += term * std::size_t{dim[N - i - 1] + LoBound};
        }
        return offset;
    }

    /**
     * @brief 根据 dim 获得对应的索引，如果索引越界则抛出异常
     *
     * @param[in] dim
     * @return constexpr std::size_t
     */
    constexpr std::size_t safe_index(const Dim &dim) const noexcept
    {
        for (std::size_t i = 0; i < N; i++)
        {
            if (dim[i] < -std::intptr_t{LoBound} || dim[i] >= m_shape[i] + HiBound)
                throw std::out_of_range("Tensor at");
        }
        return index(dim);
    }

    /**
     * @brief 获得数组首地址
     *
     * @return constexpr T*
     */
    constexpr T *data() noexcept
    {
        return m_arr.data();
    }

    /**
     * @brief 获得数组首地址
     *
     * @return constexpr const T*
     */
    constexpr const T *data() const noexcept
    {
        return m_arr.data();
    }

    /**
     * @brief 获得数组元素
     *
     * @param[in] dim
     * @return constexpr T&
     */
    constexpr T &operator()(const Dim &dim) noexcept
    {
        return data()[index(dim)];
    }

    /**
     * @brief 获得数组元素
     *
     * @param[in] dim
     * @return constexpr const T&
     */
    constexpr const T &operator()(const Dim &dim) const noexcept
    {
        return data()[index(dim)];
    }

    /**
     * @brief 多维索引元素
     *
     * @tparam Ts
     */
    template <class... Ts>
        requires(DimParams<N, Ts...>)
    constexpr T &operator()(const Ts &...ts) noexcept
    {
        return operator()(Dim{std::intptr_t(ts)...});
    }

    /**
     * @brief 多维索引元素
     *
     * @tparam Ts
     */
    template <class... Ts>
        requires(DimParams<N, Ts...>)
    constexpr const T &operator()(const Ts &...ts) const noexcept
    {
        return operator()(Dim{std::intptr_t(ts)...});
    }

    /**
     * @brief 安全访问数组元素
     *
     * @param[in] dim
     * @return T&
     */
    T &at(const Dim &dim)
    {
        return data()[safe_index(dim)];
    }

    /**
     * @brief 安全访问数组元素
     *
     * @param[in] dim
     * @return const T&
     */
    const T &at(const Dim &dim) const
    {
        return data()[safe_index(dim)];
    }

    /**
     * @brief 安全访问数组元素
     *
     * @tparam Ts
     */
    template <class... Ts>
        requires(DimParams<N, Ts...>)
    T &at(const Ts &...ts)
    {
        return at(Dim{ts...});
    }

    /**
     * @brief 安全访问数组元素
     *
     * @tparam Ts
     */
    template <class... Ts>
        requires(DimParams<N, Ts...>)
    const T &at(const Ts &...ts) const
    {
        return at(Dim{ts...});
    }

    /**
     * @brief 展平为向量，允许退化为 1 维数组访问
     *
     * @return auto
     */
    const auto &flatten() const noexcept
    {
        return m_arr;
    }

    auto &flatten() noexcept
    {
        return m_arr;
    }
};

} // namespace math

} // namespace xy
